import sys
from BasicLibraries import *
import pickle
import allocation_ga as af


#====
ONEDRIVE = ''
#====


#==== Binarisation quantile (for the robustness tests)
try:
    BQL = sys.argv[1]
    energy_control = sys.argv[2]
except Exception:
    #=== If no input is given, then go for the default binarisation threshold
    BQL=str(0.75)
#=== Default choice to control for energy prices
energy_control = False
print('Binarisation threshold:', BQL)
print('Control for energy prices for robustness tests:', energy_control)
BQL = float(BQL)

#====
AM = af.AllocationGAD()
#==== Choose your population
population_type_ = 'LandscapePopulation_0724_full.txt'
#================
np.random.seed(5)
#================

I = [        'association',
             'adoption of standards and rules',
             'assessment and measurement',
             'organizational structuring',
             'modification of procedures',
             'asset modification',
             'training',
             'r&d investments',
             'new products',
     ]
env_sdgs = ['water & energy', 'cons. & prod.', 'biodiversity']


def get_bounds(X_sorted, k_max, conf):
    lower_b_ = X_sorted.iloc[int(k_max*0.5*conf-1)]
    upper_b_ = X_sorted.iloc[int((1-0.5*conf)*k_max)]
    lower_b = np.min([lower_b_, upper_b_])
    upper_b = np.max([lower_b_, upper_b_])

    return lower_b, upper_b

#%% 
rs = []
n_iter = 1000
confidence = 0.05
for FW in [0.0, 0.25, 0.5, 0.75, 1.0]:
    print('Realisation:', FW)
    population_type = str(FW).replace('.', '_')+'_'+population_type_
    dt = pd.read_csv(population_type, low_memory=False)
    gvkey_list = list(np.unique(dt.gvkey))
    
    dat = dt.copy()
       
    rsBT = []
    for i in range(n_iter):
        datBT = dat.iloc[np.random.choice(len(dat), int(0.9*len(dat)), replace = True)].reset_index(drop = True)
        genesL, obs_fitnessL, linear_landscape, _, _, _, _,r,s = AM.fit_model(datBT,  
                                                                              sdgs_number=env_sdgs,
                                                                              dimension_ = I, \
                                                                              gvkey_list = gvkey_list,
                                                                              model_type = 'linear',
                                                                              energy_price_control = energy_control,
                                                                              normal_fit=False,
                                                                              FW= FW)  
            
        params = s.params
        params = params*genesL.std()/obs_fitnessL.std()
        r = r[r.abs() < r.abs().quantile(0.99)]
        rs_ratio = np.std(r)/params.abs().mean()
        rsBT.append(rs_ratio)


    lb,ub = np.quantile(rsBT, 0.05), np.quantile(rsBT, 0.95)
    average_value = np.mean(rsBT)
    rs.append([FW, average_value, lb, ub])
rs = pd.DataFrame(rs, columns = ['Weight', 'r/s', 'lower_bound', 'upper_bound'])
pickle.dump(rs, open('rugedness_test_rs.pckl', 'wb'))


